
import json
import logging
import os
import numpy as np

import tensorflow as tf

from flags import parse_args
from model import Model
from slim.datasets import dataset_utils

logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s', level=logging.DEBUG)

FLAGS, unparsed = parse_args()

embedding_file = FLAGS.embedding_dir + '/embeddings.npy'
vgg_checkpoint_path = FLAGS.checkpoint_path + '/vgg_16.ckpt'

with open(FLAGS.embedding_dir + '/dictionary.json', 'r') as f:
    char_to_ix = json.load(f,encoding='utf-8')

with open(FLAGS.embedding_dir + '/reversed_dictionary.json', 'r') as f:
    ix_to_char_temp = json.load(f,encoding='utf-8')

ix_to_char = {}
for key, val in ix_to_char_temp.items():
    ix_to_char[int(key)] = val

model = Model(learning_rate=FLAGS.learning_rate, batch_size=FLAGS.batch_size, state_size=FLAGS.state_size)
model.build(embedding_file = embedding_file, vgg_checkpoint_path=vgg_checkpoint_path, is_training=False)

image_filename_placeholder = tf.placeholder(dtype=tf.string)
image_tensor = tf.read_file(image_filename_placeholder)
image_tensor = tf.image.decode_jpeg(image_tensor, channels=3)
images = tf.image.resize_images(image_tensor, [model.default_image_size, model.default_image_size])
images = tf.expand_dims(images, 0)

with tf.Session() as sess:
    saver = tf.train.Saver()
    latest_ckp = tf.train.latest_checkpoint(FLAGS.output_dir)
    if latest_ckp:
        logging.debug('restore checkpoint {0}'.format(latest_ckp))
        saver.restore(sess, latest_ckp)

        logging.debug('begin running the init op')
        init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
        logging.debug('init global and local variables')
        sess.run(init_op)
        logging.debug('restoring the vgg weights')
        model.init_fn(sess)
        logging.debug('done running the init op')

        training_state = sess.run(model.init_state)
        images_val = sess.run(images, feed_dict={image_filename_placeholder:FLAGS.dataset_val})

        result = ""
        for i in range(FLAGS.seq_nums):
            has_caption = True
            feed_dict_to_use = {}
            if i == 0:
                has_caption = False
                feed_dict_to_use = {model.image_x:images_val, model.caption_x:[[]],model.init_state:training_state, model.has_caption:has_caption}
            else:
                caption_x = []
                for c in result.split(' '):
                    if c in char_to_ix:
                        caption_x.append(char_to_ix[c])
                    else:
                        caption_x.append(char_to_ix['UNK'])
                caption_x = np.array([caption_x])
                #np.array([[char_to_ix[c] for c in result.split(' ')]])
                feed_dict_to_use = {model.image_x:images_val, model.caption_x:caption_x,model.init_state:training_state, model.has_caption:has_caption}
            training_state, prediction_values = sess.run([model.final_state, model.predictions], feed_dict=feed_dict_to_use)
            char_results = np.argmax(prediction_values[0], axis=-1)
            result = ' '.join(ix_to_char[ix] for ix in char_results.ravel())
            logging.debug(result)
        logging.debug('Prediction Result:{0}'.format(result.replace('UNK', '')))

        